import yfinance as yf
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.font_manager import FontProperties

# 获取股票数据
symbol = '601919.SS'  # 股票代码
start = '2015-01-01'  # 开始时间
end = '2020-12-31'    # 结束时间
df = yf.download(symbol, start=start, end=end)

# 计算唐奇安通道
n = 20  # 计算周期
high = df['High'].rolling(n).max()
low = df['Low'].rolling(n).min()
ma = df['Close'].rolling(n).mean()

dc = pd.DataFrame()
for i in range(n-1, len(df)):
    high_value = high[i]
    low_value = low[i]
    ma_value = ma[i]
    dc = dc.append({'Date': df.index[i], 'Upper': high_value, 'Lower': low_value, '20MA': ma_value}, ignore_index=True)

# 计算和标注买卖信号
df['Buy'] = False
df['Sell'] = False
for i in range(n-1, len(df)):
    if df['Close'][i] > dc['Upper'][i-n+1:i+1].max():  # 突破上轨
        df['Buy'][i] = True
    elif df['Close'][i] < dc['Lower'][i-n+1:i+1].min():  # 突破下轨
        df['Sell'][i] = True

df['Position'] = df['Buy'].cumsum() - df['Sell'].cumsum()
df['Position'] = df['Position'].replace(to_replace=-1, method='ffill')   # 解决信号只出现在最后一天的情况

# 计算收益率和绘制收益曲线
df['Returns'] = df['Close'].pct_change()
df['Strategy'] = df['Returns'] * df['Position'].shift(1)
df[['Returns', 'Strategy']].cumsum().apply(np.exp).plot(figsize=(10, 6))
myfont = FontProperties(fname='/System/Library/Fonts/PingFang.ttc')
plt.rcParams['axes.unicode_minus'] = False
plt.xlabel('Date 日期')
plt.ylabel('Cumulative Returns 累积收益')
plt.title('Turtle Trading Strategy based on Donchian Channel')
plt.legend(loc='upper left')
plt.show()